import numpy as np
from numpy.ma import cos
import matplotlib.pyplot as plt
from matplotlib import cm
from mpl_toolkits.mplot3d import Axes3D
import math
import datetime
import matplotlib

matplotlib.use('TkAgg')  # 或者 'Qt5Agg', 'WebAgg'


class GeneticAlgorithm:
    def __init__(self):
        # 参数初始化
        self.DNA_SIZE = 24  # 编码长度
        self.POP_SIZE = 100  # 种群大小
        self.CROSS_RATE = 0.8  # 交叉率
        self.MUTA_RATE = 0.15  # 变异率
        self.Iterations = 50  # 代次数
        self.X_BOUND = [0, 10]  # X区间
        self.Y_BOUND = [0, 10]  # Y区间

        # 初始化种群
        self.pop = np.random.randint(2, size=(self.POP_SIZE, self.DNA_SIZE * 2))

        # 存储每代的最佳结果
        self.best_fitness_history = []
        self.best_individual_history = []
        self.average_fitness_history = []

        # 创建图形窗口
        self.fig = plt.figure(figsize=(16, 8))

        # 3D可视化设置
        self.ax_3d = self.fig.add_subplot(121, projection='3d')
        self.setup_3d_visualization()

        # 2D折线图设置
        self.ax_2d = self.fig.add_subplot(122)
        self.setup_2d_visualization()

        # 关闭交互模式
        plt.ioff()

    def F(self, x, y):  # 适应度函数
        return (6.452 * (x + 0.125 * y) * (cos(x) - cos(2 * y)) ** 2) / (
                    0.8 + (x - 4.2) ** 2 + 2 * (y - 7) ** 2) + 3.226 * y

    def decodeDNA(self, pop):  # 解码
        x_pop = pop[:, 1::2]  # 奇数列表示X
        y_pop = pop[:, ::2]  # 偶数列表示y
        x = x_pop.dot(2 ** np.arange(self.DNA_SIZE)[::-1]) / float(2 ** self.DNA_SIZE - 1) * (
                    self.X_BOUND[1] - self.X_BOUND[0]) + self.X_BOUND[0]
        y = y_pop.dot(2 ** np.arange(self.DNA_SIZE)[::-1]) / float(2 ** self.DNA_SIZE - 1) * (
                    self.Y_BOUND[1] - self.Y_BOUND[0]) + self.Y_BOUND[0]
        return x, y

    def get_fitness(self, pop):
        x, y = self.decodeDNA(pop)
        temp = self.F(x, y)
        return (temp - np.min(temp)) + 0.0001  # 防止适应度出现负数

    def select(self, pop, fitness):  # 根据适应度选择
        temp = np.random.choice(np.arange(self.POP_SIZE), size=self.POP_SIZE, replace=True,
                                p=(fitness) / (fitness.sum()))
        return pop[temp]

    def crossover_and_mutation(self, pop, CROSS_RATE):
        new_pop = []
        for i in pop:  # 遍历种群中的每一个个体，将该个体作为父代
            temp = i  # 子代先得到父亲的全部基因
            if np.random.rand() < CROSS_RATE:  # 以交叉概率发生交叉
                j = pop[np.random.randint(self.POP_SIZE)]  # 从种群中随机选择另一个个体作为母代
                cpoints1 = np.random.randint(0, self.DNA_SIZE * 2 - 1)  # 随机产生交叉的点
                cpoints2 = np.random.randint(cpoints1, self.DNA_SIZE * 2)
                temp[cpoints1:cpoints2] = j[cpoints1:cpoints2]  # 子代得到位于交叉点后的母代的基因
            self.mutation(temp, self.MUTA_RATE)  # 后代以变异率发生变异
            new_pop.append(temp)
        return new_pop

    def mutation(self, temp, MUTA_RATE):
        if np.random.rand() < MUTA_RATE:  # 以MUTA_RATE的概率进行变异
            mutate_point = np.random.randint(0, self.DNA_SIZE)  # 随机产生一个实数，代表要变异基因的位置
            temp[mutate_point] = temp[mutate_point] ^ 1  # 将变异点的二进制为反转

    def print_info(self, pop):  # 用于输出结果
        fitness = self.get_fitness(pop)
        maxfitness = np.argmax(fitness)  # 返回最大值的索引值
        print("max_fitness:", fitness[maxfitness])
        x, y = self.decodeDNA(pop)
        print("最优的基因型：", pop[maxfitness])
        print("(x, y):", (x[maxfitness], y[maxfitness]))
        print("F(x,y)_max = ", self.F(x[maxfitness], y[maxfitness]))

    def setup_3d_visualization(self):
        # 绘制3D表面
        X = np.linspace(*self.X_BOUND, 100)
        Y = np.linspace(*self.Y_BOUND, 100)
        X, Y = np.meshgrid(X, Y)
        Z = self.F(X, Y)
        self.ax_3d.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.coolwarm, alpha=0.6)
        self.ax_3d.set_zlim(-20, 100)
        self.ax_3d.set_xlabel('x')
        self.ax_3d.set_ylabel('y')
        self.ax_3d.set_zlabel('z')
        self.ax_3d.set_title('3D Function Visualization with Population')

    def setup_2d_visualization(self):
        # 设置2D折线图
        self.ax_2d.set_xlim(0, self.Iterations)
        self.ax_2d.set_ylim(0, 100)  # 假设适应度最大值为100
        self.ax_2d.set_xlabel('Generation')
        self.ax_2d.set_ylabel('Fitness')
        self.ax_2d.set_title('Fitness Trend Over Generations')
        self.ax_2d.grid(True)

        # 初始化线条
        self.best_line, = self.ax_2d.plot([], [], 'r-', label='Best Fitness')
        self.avg_line, = self.ax_2d.plot([], [], 'b-', label='Average Fitness')
        self.ax_2d.legend()

    def update_visualizations(self, pop, generation):
        x, y = self.decodeDNA(pop)
        fitness = self.get_fitness(pop)

        # 更新3D可视化
        if hasattr(self, 'sca_3d'):
            self.sca_3d.remove()
        self.sca_3d = self.ax_3d.scatter(x, y, self.F(x, y), c='red', marker='o', s=50, alpha=0.8)

        # 记录当前代的统计数据
        self.best_fitness_history.append(np.max(fitness))
        self.average_fitness_history.append(np.mean(fitness))

        # 更新2D折线图
        generations = range(generation + 1)
        self.best_line.set_data(generations, self.best_fitness_history)
        self.avg_line.set_data(generations, self.average_fitness_history)

        # 调整Y轴范围以适应数据
        current_max = max(self.best_fitness_history)
        self.ax_2d.set_ylim(0, current_max * 1.1)

        # 重绘图形
        plt.draw()
        plt.pause(0.05)

    def run(self):
        print('==========迭代=============')

        for generation in range(self.Iterations):  # 迭代N代
            print(f"当前代数: {generation}/{self.Iterations}")
            x, y = self.decodeDNA(self.pop)

            # 获取当前代的适应度
            fitness = self.get_fitness(self.pop)
            maxfitness = np.max(fitness)

            # 获取当前代的最优个体
            best_idx = np.argmax(fitness)
            x, y = self.decodeDNA(self.pop)
            self.best_individual_history.append((x[best_idx], y[best_idx]))

            # 更新可视化
            self.update_visualizations(self.pop, generation)

            # 进行遗传操作
            self.pop = np.array(self.crossover_and_mutation(self.pop, self.CROSS_RATE))
            self.pop = self.select(self.pop, fitness)  # 选择生成新的种群

        # 输出最终结果
        self.print_info(self.pop)
        plt.show()  # 保持图形窗口打开


if __name__ == "__main__":
    ga = GeneticAlgorithm()
    ga.run()